from typing import (List, Tuple, Union, Dict)
from collections import namedtuple
import pandas as pd
import numpy as np

import json
# import seaborn as sns
# import plotly.express as px
# import plotly.graph_objects as go
# from plotly.subplots import make_subplots

# import matplotlib.pyplot as plt
# plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
# plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

__all__ = [
    'pct_chg',
    'top_n2res',
    'show_rs',
    'top_n2wres',
    'show_wrs',
    'low5_bar_list',
    'top5_concept_price',
]


def calc_rs(df: pd.DataFrame, col: str) -> pd.DataFrame:
    """计算相对强弱指标 RS

    Args:
        df (pd.DataFrame): index-date columns-symbol value
        col (int): 窗口期

    Raises:
        ValueError: 窗口不能大于df

    Returns:
        pd.DataFrame: N日RS
    """
    if col not in df.columns:
        raise ValueError('col不在df中')

    pct_chg = df[col]  # N日收益率
    # 横向排序 降序 归一化
    rank = pct_chg.rank(pct=True, ascending=True)

    return rank


def calc_pct_nd(price: pd.DataFrame, nd: Union[int, List]) -> pd.DataFrame:
    """计算n_days的收益率情况

    Parameters
    ----------
    price : pd.DataFrame
        _description_
    nd : Union[int,List]
        _description_

    Returns
    -------
    pd.DataFrame
        _description_
    """
    if isinstance(nd, int):

        nd = [nd]

    cols = [f"{i}日" for i in nd]

    pct_chg: pd.DataFrame = pd.concat(
        (price.pct_change(i).iloc[-1] for i in nd), axis=1)
    pct_chg.columns = cols

    return pct_chg


def get_topn2cons(cons_name: str,
                  rank_col: str,
                  concept_dic: Dict,
                  price: pd.DataFrame,
                  cols: List = [1, 5, 10, 20, 60, 120]) -> pd.DataFrame:
    """获取指定概念 N日动量个股情况

    Parameters
    ----------
    cons_name : str
        指定概念
    rank_col : str
        按N日排名
    concept_dic : Dict
        概念字典
    price : pd.DataFrame
        价格数据

    Returns
    -------
    pd.DataFrame
    """
    target = concept_dic[0].get(cons_name, '')
    if target != '':
        target = [concept_dic[1][i] for i in target]

        cons_pct = calc_pct_nd(price, cols)
        return cons_pct.loc[target, :].sort_values(rank_col, ascending=False)


def calc_rs_gn(df: pd.DataFrame, N: int) -> pd.DataFrame:
    """计算相对强弱指标 RS

    Args:
        df (pd.DataFrame): index-date columns-symbol value
        N (int): 窗口期

    Raises:
        ValueError: 窗口不能大于df

    Returns:
        pd.DataFrame: N日RS
    """
    if len(df) <= N:
        raise ValueError('参数N不能大于df的长度')

    pct_chg = df.pct_change(N)  # N日收益率
    # 横向排序 降序 归一化
    rank = pct_chg.rank(axis=1, pct=True, ascending=True)

    return rank


"""同花顺概念部分"""

# 获取同花顺
data = pd.read_csv(r'data\ths_data.csv',
                   encoding='gbk',
                   index_col=[0],
                   parse_dates=['日期'])

# 动量前五概念成分股数据获取
with open(r'Data\concept_dic.json', 'r') as file:
    concept_dic = json.loads(file.read())

# 转换数据结构
pivot_df = pd.pivot_table(data, index='日期', columns='概念名称', values='收盘价')

# 计算涨跌幅
pct_chg: pd.DataFrame = calc_pct_nd(pivot_df, [5, 10, 20, 60, 120])

# 动量得分-等权
rs_df = calc_rs_gn(pivot_df, 20) + calc_rs_gn(pivot_df, 40) + calc_rs_gn(
    pivot_df, 60)
rs_df /= 3

top_n2res = rs_df.iloc[-1].nlargest(20).index.tolist()

show_rs = rs_df[top_n2res].iloc[-20:]

# 动量得分-加权
wrs_df = calc_rs_gn(pivot_df, 1)*0.3 + calc_rs_gn(pivot_df, 5)*0.4 + \
    calc_rs_gn(pivot_df, 20)*0.3
top_n2wres = wrs_df.iloc[-1].nlargest(20).index.tolist()
show_wrs = wrs_df[top_n2wres].iloc[-20:]

# 前五BigNumber
top5_bar_list = wrs_df.iloc[-1].nlargest(5).index.tolist()

top5_bar: pd.DataFrame = wrs_df[top5_bar_list].iloc[-1].to_frame('动量得分')
top5_bar['5日动量'] = pct_chg.loc[top5_bar_list, '5日']

# 后五BigNumber
low5_bar_list = wrs_df.iloc[-1].nsmallest(5).index.tolist()

low5_bar: pd.DataFrame = wrs_df[low5_bar_list].iloc[-1].to_frame('动量得分')
low5_bar['5日动量'] = pct_chg.loc[low5_bar_list, '5日']

# 概念前五名的成分股 table
top5_concept_price = pd.read_csv(r'Data\top5_concept_price.csv',
                                 index_col=[0],
                                 encoding='gbk')
